{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using a C Heartbeat function\n",
    "The heartbeat function of a REBOUND simulation gets called after every timestep. There are many different things you can do in a heartbeat function, for example creating outputs, adding particles, adjusting parameters that depend on time, etc. REBOUND supports heartbeat functions in both its C and python interface. \n",
    "\n",
    "A python heartbeat function can sometimes become the bottleneck of a simulation because it gets called every single timestep. This tutorial shows you how to implement the heartbeat function in C, then link it to REBOUND using python. Note that alternatively you can of course always just use the C version of REBOUND directly and never bother with python at all.\n",
    "\n",
    "We start by creating a REBOUND simulation which contains the planets of our Solar System as a test case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import rebound\n",
    "sim = rebound.Simulation()\n",
    "rebound.data.add_solar_system(sim)\n",
    "sim.integrator = \"whfast\"\n",
    "sim.dt = sim.particles[1].P/30.13 # About 30 steps for each Mercury Orbit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us first create a simple heartbeat function in python. It simply calculates the eccentricity of Mercury (you could do something with it, here we just calculate it and then ignore it)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Eccentricity: 0.205636 \n"
     ]
    }
   ],
   "source": [
    "# We use a global variable to store the value of the eccentricity\n",
    "e = 0 \n",
    "def heartbeat(sim_pointer):\n",
    "    global e\n",
    "    # The function argument is a pointer to the simulation:\n",
    "    # Here we get its contents:\n",
    "    sim = sim_pointer.contents \n",
    "    e = sim.particles[1].e\n",
    "sim.heartbeat = heartbeat\n",
    "sim.integrate(sim.t+1)\n",
    "print(\"Eccentricity: %f \" %e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's measure how long it takes to integrate 1000 orbits:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Runtime: 0.474900 s\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "start = time.time()\n",
    "sim.integrate(sim.t + sim.particles[1].P*1000)\n",
    "stop = time.time()\n",
    "print(\"Runtime: %f s\"%(stop-start))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now implement this in C. For this to work, it is best to download and work with a full REBOUND repository (download the package from github, rather than just installing the python package with pip install). \n",
    "\n",
    "We first write our heartbeat function in C. The following cell writes to a new file in the current directory, `heartbeat.c` (you can also use an external editor and terminal window to do the same without the jupyter magic commands):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting heartbeat.c\n"
     ]
    }
   ],
   "source": [
    "%%writefile heartbeat.c\n",
    "#include \"rebound.h\"\n",
    "double e =0; // global variable\n",
    "void heartbeat(struct reb_simulation* sim_pointer){\n",
    "    struct reb_orbit orbit = reb_orbit_from_particle(sim_pointer->G, sim_pointer->particles[0], sim_pointer->particles[1]);\n",
    "    e = orbit.e;\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we compile and link our heartbeat function as a shared library, we need the REBOUND header file and the shared library file. Different operating systems and compilers handles the paths to shared libraries differently. This can quickly get rather frustrating. If you're familiar with C, by all means go ahead and do it the proper way. A hack to get around most of these difficulties is to simply copy the REBOUND header and library to the current folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cp ../src/librebound.so .\n",
    "!cp ../src/rebound.h ."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you installed REBOUND with pip, you can look up the paths to the two files in python and e.g. create symlinks to them into your working directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/path/to/your/venv/lib/python3.9/site-packages/rebound/../librebound.cpython-39-x86_64-linux-gnu.so\n",
      "/path/to/your/venv/lib/python3.9/site-packages/rebound/rebound.h\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "print(rebound.__libpath__)\n",
    "print(Path(rebound.__file__).parent / \"rebound.h\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we can compile and link the code. `-fPIC` instructs gcc to create Position Independent Code, which might not be needed on your operating system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "!gcc -c -O3 -fPIC heartbeat.c -o heartbeat.o"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "!gcc -L. -shared heartbeat.o -o heartbeat.so -lrebound "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we load the library using ctypes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ctypes import cdll\n",
    "clibheartbeat = cdll.LoadLibrary(\"heartbeat.so\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now finally set the function pointer in our simulation to the new heartbeat function and then run the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.heartbeat = clibheartbeat.heartbeat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Runtime: 0.153589 s\n"
     ]
    }
   ],
   "source": [
    "start = time.time()\n",
    "sim.integrate(sim.t + sim.particles[1].P*1000)\n",
    "stop = time.time()\n",
    "print(\"Runtime: %f s\"%(stop-start))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the simulation runs significantly faster using the C heartbeat function as we avoid all the python overhead.\n",
    "\n",
    "We can print out the value of the global variable `e` in the heartbeat library (using global variables in a shared library is not the best way to store data - all simulations will see the same variable and you could end up with unexpected behaviour if you are running multiple simulations in one python program)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.2057293010348337\n"
     ]
    }
   ],
   "source": [
    "from ctypes import c_double\n",
    "print(c_double.in_dll(clibheartbeat,\"e\").value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
